#ifndef CODEGENERATOR_H
#define CODEGENERATOR_H

#include <string>
#include <vector>
#include <map>
#include <stack>
#include <memory>
#include <unordered_map>
#include <set>
#include "ASTNode.h"
#include "SymbolTable.h"
#include "MatrixOptimizer.h"

// RISC-V寄存器定义
enum class Register {
    ZERO,   // x0 - 零寄存器
    RA,     // x1 - 返回地址
    SP,     // x2 - 栈指针
    GP,     // x3 - 全局指针
    TP,     // x4 - 线程指针
    T0,     // x5 - 临时寄存器
    T1,     // x6 - 临时寄存器
    T2,     // x7 - 临时寄存器
    S0,     // x8 - 保存寄存器
    S1,     // x9 - 保存寄存器
    A0,     // x10 - 参数/返回值寄存器
    A1,     // x11 - 参数寄存器
    A2,     // x12 - 参数寄存器
    A3,     // x13 - 参数寄存器
    A4,     // x14 - 参数寄存器
    A5,     // x15 - 参数寄存器
    A6,     // x16 - 参数寄存器
    A7,     // x17 - 参数寄存器
    S2,     // x18 - 保存寄存器
    S3,     // x19 - 保存寄存器
    S4,     // x20 - 保存寄存器
    S5,     // x21 - 保存寄存器
    S6,     // x22 - 保存寄存器
    S7,     // x23 - 保存寄存器
    S8,     // x24 - 保存寄存器
    S9,     // x25 - 保存寄存器
    S10,    // x26 - 保存寄存器
    S11,    // x27 - 保存寄存器
    T3,     // x28 - 临时寄存器
    T4,     // x29 - 临时寄存器
    T5,     // x30 - 临时寄存器
    T6      // x31 - 临时寄存器
};

// 寄存器分配器
class RegisterAllocator {
private:
    std::map<Register, bool> registerUsage;
    std::map<std::string, Register> variableRegisters;
    std::stack<Register> freeRegisters;
    std::map<Register, int> spilledRegisters; // 溢出寄存器到栈偏移的映射
    int nextSpillOffset = -4; // 下一个溢出偏移量

public:
    RegisterAllocator();
    Register allocateRegister();
    void freeRegister(Register reg);
    void allocateSpecificRegister(Register reg);
    bool hasFreeRegisters() const;
    Register getVariableRegister(const std::string& varName);
    void setVariableRegister(const std::string& varName, Register reg);
    void clearVariableRegisters();
    std::string registerToString(Register reg);

    // 新增：寄存器溢出管理
    Register spillRegister(Register reg);
    void unspillRegister(Register reg);
    bool isCallerSaved(Register reg);
    bool isCalleeSaved(Register reg);
    std::vector<Register> getActiveCallerSavedRegisters();
    void generateSpillInstruction(Register reg, std::vector<std::string>& code);
    void generateUnspillInstruction(Register reg, std::vector<std::string>& code);
    bool isRegisterAllocated(Register reg);
    void reset();
    void resetSpillOffset();
};

// 参数信息结构体
struct ParamInfo {
    bool inRegister; // 是否在寄存器
    int regIndex;   // 如果在寄存器，a0=0, a1=1, ...
    int stackOffset; // 如果在栈，偏移量
    int index;      // 参数在形参列表中的顺序
};

// Offset类型说明：正数=寄存器编号（xi），负数=内存偏移
// 形参/临时变量区：负偏移(s0)，临时数据区：负偏移(sp)

class CodeGenerator {
private:
    std::vector<std::string> code;
    RegisterAllocator regAlloc;
    std::string currentFunction;
    std::unordered_map<std::string, int> globalVarOffsets;
    int nextOffset = 0;
    int minNextOffset = 0;
    std::unordered_map<std::string, int> allVarOffsets;
    std::vector<std::string> currentFunctionParams;
    std::unordered_map<std::string, std::pair<int, int>> currentFunctionParamInfo;
    std::vector<std::string> currentFunctionParamOrder;
    int currentFrameSize = 0;
    int tmpPtr = 0;
    int maxTempDepth = 0;
    int minVarOffset = 0;
    std::stack<std::string> continueLabels;
    std::stack<std::string> breakLabels;
    int labelCounter = 0;

    // 作用域管理相关
    struct Scope {
        int baseOffset;
        int nextOffset;
        std::unordered_map<std::string, int> varOffsets;
        int scopeId; // 作用域ID
        Scope(int baseOffset = 0) : baseOffset(baseOffset), nextOffset(baseOffset), scopeId(-1) {}
    };
    std::stack<Scope> scopeStack;
    int scopeIdCounter = 0; // 全局作用域ID计数器
    std::stack<int> baseOffsetStack; // BaseOffset栈管理

    // 全局变量offset映射数据结构：变量名@作用域ID -> offset
    std::unordered_map<std::string, int> globalVarOffsetMap; // varName@scopeId -> offset
    std::unordered_map<int, int> parentScopeMap; // scopeId -> parentScopeId

    // 新增：函数名到根作用域ID的映射表
    std::unordered_map<std::string, int> functionRootScopeIdMap; // funcName -> rootScopeId

    // 新增：CodeGen阶段的作用域跟踪
    int currentScopeId = -1; // 当前作用域ID
    std::unordered_map<int, bool> visitedScopes; // 记录已访问的作用域
    std::stack<int> codegenScopeStack; // 代码生成阶段的作用域栈，用于跟踪嵌套作用域

    // 寄存器保存相关
    std::vector<std::pair<Register, int>> savedCallerRegisters;
    int currentCallerSaveStackOffset = 0;

    // 新增：保存区最大正偏移
    int saveAreaMaxOffset = 48;
    // 新增：形参区起始负偏移
    int paramAreaStartOffset = 0;
    // 新增：临时变量区起始负偏移
    int localVarAreaStartOffset = 0;
    // 新增：临时数据区分配栈（存储已分配的sp正偏移）
    std::vector<int> tempDataOffsetStack;
    // 新增：当前可用t寄存器编号
    std::vector<int> freeTRegs;
    // 新增：当前可用s寄存器编号
    std::vector<int> freeSRegs;

    // 辅助函数
    std::string generateLabel(const std::string& prefix);
    void emit(const std::string& line);
    void emitComment(const std::string& comment);
    void emitFunctionPrologue(const std::string& funcName, int frameSize);
    void emitFunctionEpilogue(int frameSize);
    void saveRegisters();
    void restoreRegisters();

    // 表达式代码生成
    Register generateExpression(ASTNode* node, Register targetReg = Register::T0,bool tag=true);
    Register generateBinaryOp(ASTNode* node);
    Register generateUnaryOp(ASTNode* node);
    Register generateVarRef(ASTNode* node, Register targetReg,bool tag=true);
    Register generateNumLiteral(ASTNode* node);
    Register generateFuncCall(ASTNode* node);

    // 语句代码生成
    void generateStatement(ASTNode* node, bool& hasReturn, int frameSize, const std::set<std::string>& paramNameSet);
    void generateBlock(ASTNode* node, bool& hasReturn, int frameSize, const std::set<std::string>& paramNameSet);
    void generateDeclAssignStmt(ASTNode* node, const std::set<std::string>& paramNameSet);
    void generateAssignStmt(ASTNode* node);
    void generateIfStmt(ASTNode* node, bool& hasReturn, int frameSize, const std::set<std::string>& paramNameSet);
    void generateIfElseStmt(ASTNode* node, bool& hasReturn, int frameSize, const std::set<std::string>& paramNameSet);
    void generateWhileStmt(ASTNode* node, bool& hasReturn, int frameSize, const std::set<std::string>& paramNameSet);
    void generateBreakStmt(ASTNode* node);
    void generateContinueStmt(ASTNode* node);
    void generateReturnStmt(ASTNode* node, bool& hasReturn, int frameSize);

    // 条件代码生成
    void generateCondition(ASTNode* node, const std::string& trueLabel, const std::string& falseLabel);
    void generateShortCircuit(ASTNode* node, const std::string& trueLabel, const std::string& falseLabel);

    // 栈管理
    void pushToStack(Register reg);
    Register popFromStack();
    void pushTemp(Register reg);
    Register popTemp();
    int getVariableOffset(const std::string& varName);
    int findVarOffset(const std::string& varName, int scopeId);

    // 变量管理
    int countLocalVars(ASTNode* node);
    void preallocateVariables(ASTNode* node);
    int getMinVarOffset();
    void collectVarOffsets(ASTNode* node, const std::set<std::string>& paramNameSet);
    void rebuildScopeStack(ASTNode* node, std::stack<Scope>& tempScopeStack, const std::set<std::string>& paramNameSet);

    // 作用域管理函数
    void enterScope(); // 进入作用域
    void enterScope(ASTNode* node); // 带AST节点的进入作用域
    void exitScope(); // 退出作用域
    int findVariableOffset(const std::string& varName, int scopeId); // 查找变量offset
    void allocateVar(const std::string& varName); // 分配变量
    void assignOffsets(ASTNode* node, int currentScopeId, int parentScopeId);
    bool isVariableInCurrentScope(const std::string& varName);

    // 新增：全局变量offset映射管理
    void addVarOffset(const std::string& varName, int scopeId, int offset); // 添加变量offset映射
    int findVarOffsetByScope(const std::string& varName, int scopeId); // 根据作用域查找变量offset
    std::string makeVarKey(const std::string& varName, int scopeId); // 生成变量键值

    // 新增：CodeGen阶段的作用域管理函数
    void enterFunctionScope(const std::string& funcName); // 进入函数作用域
    void enterBlockScope(); // 进入代码块作用域
    void exitBlockScope(); // 退出代码块作用域
    int getCurrentScopeId() const; // 获取当前作用域ID
    void resetScopeTracking(); // 重置作用域跟踪状态

    // 寄存器保存相关
    void saveCallerRegisters();
    void restoreCallerRegisters();
    void generateSpillInstructions();
    void generateUnspillInstructions();

    // Offset分配辅助函数
    int allocateParamOffset(); // 分配形参区offset
    int allocateLocalVarOffset(); // 分配临时变量区offset
    int allocateTempDataOffset(); // 分配临时数据区offset
    int allocateTReg(); // 分配t寄存器
    int allocateSReg(); // 分配s寄存器
    void freeTReg(int reg); // 回收t寄存器
    void freeSReg(int reg); // 回收s寄存器

    // Offset->寄存器/内存映射辅助
    std::string offsetToReg(int offset); // 正数offset转寄存器名
    std::string offsetToMem(int offset, bool isS0Base); // 负数offset转内存地址

    // 新增：分段映射函数
    int getVarOffsetByPtr(int ptr);
    int getDataOffsetByPtr(int ptr);
    // 新增：重载assignOffsets
    void assignOffsets(ASTNode* node, int currentScopeId, int parentScopeId, int& NextVarPtr, std::stack<int>& BaseVarPtrStack, int& NextDataPtr);

    // 新增：多参数传递相关全局参数
    bool ParaOverflow = false;  // 是否正在处理多余参数
    int BaseDataOffset = 0;     // 临时数据偏移量调整

    // 新增：矩陣優化器
    MatrixOptimizer matrixOptimizer;
    
    // 新增：性能統計
    struct {
        double totalOptimizationGain = 0.0;
    } stats;

    // 新增：矩陣優化與代碼生成集成函數
    bool isMatrixOperation(ASTNode* node);
    Register generateMatrixExpression(ASTNode* node, Register targetReg);
    Register generateOptimizedMatrixAddition(ASTNode* node, Register targetReg);
    Register generateOptimizedMatrixMultiplication(ASTNode* node, Register targetReg);
    Register generateVectorizedMatrixAddition(ASTNode* left, ASTNode* right, Register targetReg);
    Register generateParallelMatrixAddition(ASTNode* left, ASTNode* right, Register targetReg);
    Register generateStrassenMatrixMultiplication(ASTNode* left, ASTNode* right, Register targetReg);
    Register generateBlockMatrixMultiplication(ASTNode* left, ASTNode* right, Register targetReg);
    Register generateVectorizedMatrixMultiplication(ASTNode* left, ASTNode* right, Register targetReg);
    Register generateStandardMatrixAddition(ASTNode* left, ASTNode* right, Register targetReg);
    Register generateStandardMatrixMultiplication(ASTNode* left, ASTNode* right, Register targetReg);
    Register generateOptimizedMatrixFunction(ASTNode* node, Register targetReg);
    Register generateOptimizedMatrixTranspose(ASTNode* node, Register targetReg);
    Register generateOptimizedMatrixInverse(ASTNode* node, Register targetReg);
    Register generateOptimizedMatrixDeterminant(ASTNode* node, Register targetReg);

public:
    CodeGenerator();
    ~CodeGenerator();
    void generate(ASTNode* root);
    void printCode() const;
    void saveToFile(const std::string& filename) const;
    void optimizeCommonSubexpressions();
    void optimizeConstantFolding();
    void optimizeRegisterAllocation();
    void optimizeMatrixOperations();
    int calcMaxTempDepth(ASTNode* node, int curDepth);
};

#endif // CODEGENERATOR_H
